import os
from wechat import Wechat
from wechat_db import WechatDB

if __name__ == "__main__":
    wx = Wechat()
    wx_aes_key = wx.wechat_init()
    version = wx.get_wechat_version()
    base = hex(wx.get_wechat_base())
    aes_key = wx_aes_key.get_aes_key()

    print(f"当前微信版本号 {version}")
    print(f"当前微信基址 {base}")

    if aes_key is not None:
        print(f"当前微信AES密钥 {aes_key}")
        db = WechatDB(aes_key)
        try:
            user_id = None
            db_folders = []

            # 获取微信目录
            user_directory = db.get_wechat_user_directory()

            # 获取目录下全部wxid_开头的文件夹
            folders = [
                folder
                for folder in os.listdir(user_directory)
                if folder.startswith("wxid_")
            ]

            # 筛选出wxid_目录里面有Msg目录的文件夹
            msg_folders = [
                folder
                for folder in folders
                if os.path.isdir(os.path.join(user_directory, folder, "Msg"))
            ]

            # 对文件夹的db进行解密测试
            for folder in msg_folders:
                # 完整路径
                folder_path = os.path.join(user_directory, folder, "Msg")
                try:
                    # 用ChatMsg.db进行解密测试
                    header, key, pages = db.decrypt_db(
                        os.path.join(folder_path, "ChatMsg.db")
                    )

                    user_id = folder
                    print(f"当前微信号 {user_id}")

                    # 递归找出所有需要解密的db文件
                    for root, dirs, files in os.walk(folder_path):
                        for file in files:
                            if file.endswith(".db"):
                                db_folders.append(os.path.join(root, file))

                except Exception as e:
                    print(f"解密失败 {folder} {e}")

            # 解密所有db文件
            for db_folder in db_folders:
                file_name = os.path.basename(db_folder)
                try:
                    save_path_directory = os.path.join(
                        db.save_path_directory, user_id
                    )  # 保存目录
                    header, key, pages = db.decrypt_db(db_folder)  # 开始解密

                    # 如果用户目录不存在
                    if not os.path.exists(save_path_directory):
                        os.makedirs(save_path_directory)

                    # 解密成功保存db文件
                    save_path = os.path.join(save_path_directory, file_name)
                    db.save_db(
                        header,
                        key,
                        pages,
                        save_path,
                    )
                    print(f"解密成功 {file_name} 保存路径: {save_path}")

                except Exception as e:
                    print(f"解密失败 {file_name} {e}")
        except Exception as e:
            print(e)
    else:
        print("未获取到微信AES密钥，请确认微信已经登陆")
